"""
球化组织分类分析服务

提供文件夹级别的球化组织分类分析功能。
"""

import os
import sys
import time
from typing import Dict, Any, List
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))
sys.path.insert(0, project_root)

from devdeploy.src import Classifier, get_optimal_device, prepare_output_directory


class BallLevelService:
    """
    球化组织分类分析服务
    
    提供文件夹级别的球化组织分类分析功能。
    """
    
    def __init__(self, model_path: str, device: str = None):
        """
        初始化球化组织分类分析服务
        
        Args:
            model_path (str): 模型文件路径
            device (str, optional): 推理设备，如果为None则自动选择
        """
        self.model_path = model_path
        self.device = device or get_optimal_device()
        
        # 初始化分类器
        try:
            self.classifier = Classifier(
                model_path=self.model_path,
                device=self.device
            )
            print(f"球化组织分类分析服务初始化成功，使用设备: {self.device}")
        except Exception as e:
            raise RuntimeError(f"初始化球化组织分类分析服务失败: {str(e)}")
    
    def analyze_folder(self, input_folder: str, output_folder: str) -> Dict[str, Any]:
        """
        分析文件夹中的所有图像
        
        Args:
            input_folder (str): 输入文件夹路径
            output_folder (str): 输出文件夹路径
            
        Returns:
            Dict[str, Any]: 分析结果
        """
        result = {
            'success': False,
            'reuslts': [],
            'message': '',
            'total_images': 0,
            'processed_images': 0,
            'failed_images': 0
        }
        
        try:
            # 验证输入文件夹
            if not os.path.exists(input_folder):
                result['message'] = f"输入文件夹不存在: {input_folder}"
                return result
            
            if not os.path.isdir(input_folder):
                result['message'] = f"路径不是文件夹: {input_folder}"
                return result
            
            # 使用prepare_output_directory确保输出文件夹正确创建
            if not prepare_output_directory(output_folder, clear_existing=True):
                result['message'] = f"无法创建输出文件夹: {output_folder}"
                return result
            
            # 获取图像文件列表
            image_files = self._get_image_files(input_folder)
            result['total_images'] = len(image_files)
            
            if len(image_files) == 0:
                result['message'] = f"文件夹中没有找到支持的图像文件: {input_folder}"
                return result
            
            print(f"开始分析 {len(image_files)} 个图像文件...")
            
            # 存储所有分析结果用于生成报告
            all_results = []
            
            # 逐个处理图像
            for i, image_path in enumerate(image_files):
                try:
                    print(f"处理图像 {i+1}/{len(image_files)}: {os.path.basename(image_path)}")
                    
                    # 分类单张图像
                    start_time = time.time()
                    classification_result = self.classifier.classify_single(image_path)
                    elapsed_time = time.time() - start_time
                    print(f"  - 单图像分类耗时: {elapsed_time:.3f} 秒")
                    
                    if not classification_result.get('success', False):
                        print(f"  - 分类失败: {classification_result.get('error', '未知错误')}")
                        result['failed_images'] += 1
                        continue
                    
                    # 获取文件名（不含扩展名）
                    file_name = Path(image_path).stem
                    
                    # 记录结果
                    all_results.append({
                        'file_name': file_name,
                        'class_id': classification_result.get('class_id'),
                        'class_name': classification_result.get('class_name', ''),
                        'confidence': classification_result.get('confidence', 0.0),
                        # 'probabilities': classification_result.get('probabilities', [])
                    })
                    
                    print(f"  - 分类结果: {classification_result.get('class_name')} (ID: {classification_result.get('class_id')}, 置信度: {classification_result.get('confidence', 0):.3f})")
                    
                    # 绘制结果到图像并保存
                    output_image_path = os.path.join(output_folder, f"{file_name}_result.jpg")
                    self._draw_result_on_image(image_path, classification_result, output_image_path)
                    
                    result['processed_images'] += 1
                    
                except Exception as e:
                    print(f"处理图像失败 {image_path}: {str(e)}")
                    result['failed_images'] += 1
            
            # 生成汇总报告
            self._generate_summary_report(all_results, output_folder)
            
            result['success'] = True
            result['results'] = all_results
            result['message'] = f"分析完成！成功处理 {result['processed_images']} 个文件，失败 {result['failed_images']} 个文件"
            
        except Exception as e:
            result['message'] = f"分析文件夹时发生错误: {str(e)}"
        
        return result
    
    def _draw_result_on_image(self, image_path: str, classification_result: Dict[str, Any], output_path: str) -> bool:
        """
        在图像左上角绘制分类结果并保存
        
        使用PIL进行绘制，支持中文，字体大小根据图像尺寸自动调整。
        
        Args:
            image_path: 原始图像路径
            classification_result: 分类结果
            output_path: 输出图像路径
            
        Returns:
            bool: 是否保存成功
        """
        try:
            # 使用PIL读取图像
            img = Image.open(image_path).convert('RGB')
            img_width, img_height = img.size
            
            # 获取分类结果信息
            class_name = classification_result.get('class_name', '未知')
            class_id = classification_result.get('class_id', '')
            confidence = classification_result.get('confidence', 0.0)
            
            # 准备文本
            text_lines = [
                f"类别: {class_name}",
                f"ID: {class_id}",
                f"置信度: {confidence:.3f}"
            ]
            
            # 根据图像尺寸计算字体大小（图像宽度的2-3%）
            base_font_size = max(16, int(min(img_width, img_height) * 0.025))
            
            # 尝试加载中文字体
            font = None
            font_paths = [
                # macOS 中文字体
                "/System/Library/Fonts/PingFang.ttc",
                "/System/Library/Fonts/STHeiti Light.ttc",
                "/System/Library/Fonts/Helvetica.ttc",
                # Windows 中文字体
                "C:/Windows/Fonts/simhei.ttf",  # 黑体
                "C:/Windows/Fonts/simsun.ttc",  # 宋体
                "C:/Windows/Fonts/msyh.ttc",    # 微软雅黑
                # Linux 中文字体
                "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",
                "/usr/share/fonts/truetype/arphic/uming.ttc",
            ]
            
            for font_path in font_paths:
                try:
                    if os.path.exists(font_path):
                        # 对于.ttc文件，需要指定字体索引
                        if font_path.endswith('.ttc'):
                            font = ImageFont.truetype(font_path, base_font_size, index=0)
                        else:
                            font = ImageFont.truetype(font_path, base_font_size)
                        break
                except:
                    continue
            
            # 如果找不到字体，使用默认字体
            if font is None:
                try:
                    font = ImageFont.load_default()
                except:
                    font = None
            
            # 创建绘图对象
            draw = ImageDraw.Draw(img)
            
            # 设置颜色
            text_color = (0, 255, 0)  # 绿色 (RGB)
            bg_color = (0, 0, 0)  # 黑色背景
            bg_alpha = 180  # 背景透明度 (0-255)
            
            # 计算文本尺寸和位置
            padding = max(10, int(min(img_width, img_height) * 0.01))
            line_spacing = max(5, int(base_font_size * 0.3))
            
            # 计算每行文本的尺寸
            line_heights = []
            max_text_width = 0
            
            for line in text_lines:
                if font:
                    bbox = draw.textbbox((0, 0), line, font=font)
                    text_width = bbox[2] - bbox[0]
                    text_height = bbox[3] - bbox[1]
                else:
                    # 使用默认字体时的估算
                    text_width = len(line) * base_font_size * 0.6
                    text_height = base_font_size
                
                line_heights.append(text_height)
                max_text_width = max(max_text_width, text_width)
            
            # 计算背景框尺寸
            total_height = sum(line_heights) + line_spacing * (len(text_lines) - 1) + padding * 2
            bg_width = max_text_width + padding * 2
            
            # 创建半透明背景
            bg_overlay = Image.new('RGBA', (bg_width, total_height), (*bg_color, bg_alpha))
            img.paste(bg_overlay, (padding, padding), bg_overlay)
            
            # 绘制文本
            y_pos = padding
            for i, line in enumerate(text_lines):
                x_pos = padding + 5
                draw.text((x_pos, y_pos), line, fill=text_color, font=font)
                y_pos += line_heights[i] + line_spacing
            
            # 保存图像
            img.save(output_path, quality=95)
            print(f"  - 结果图像已保存: {output_path}")
            return True
            
        except Exception as e:
            print(f"  - 绘制结果图像失败: {str(e)}")
            return False
    
    def _get_image_files(self, folder_path: str) -> List[str]:
        """
        获取文件夹中的所有图像文件
        
        Args:
            folder_path (str): 文件夹路径
            
        Returns:
            List[str]: 图像文件路径列表
        """
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
        image_files = []
        
        for file_path in Path(folder_path).iterdir():
            if file_path.is_file() and file_path.suffix.lower() in image_extensions:
                image_files.append(str(file_path))
        
        return sorted(image_files)
    
    def _generate_summary_report(self, all_results: List[Dict], output_folder: str):
        """
        生成汇总报告到根目录的result.txt文件
        
        Args:
            all_results (List[Dict]): 所有分析结果
            output_folder (str): 输出文件夹
        """
        try:
            summary_path = os.path.join(output_folder, 'result.txt')
            
            with open(summary_path, 'w', encoding='utf-8') as f:
                f.write("球化组织分类分析汇总报告\n")
                f.write("=" * 50 + "\n")
                f.write(f"分析时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"模型路径: {self.model_path}\n")
                f.write(f"总文件数: {len(all_results)}\n\n")
                
                if all_results:
                    # 统计各类别的数量和平均置信度
                    class_stats = {}
                    for r in all_results:
                        class_id = r.get('class_id')
                        class_name = r.get('class_name', f'类别{class_id}')
                        confidence = r.get('confidence', 0.0)
                        
                        if class_id not in class_stats:
                            class_stats[class_id] = {
                                'class_name': class_name,
                                'count': 0,
                                'total_confidence': 0.0,
                                'confidences': []
                            }
                        
                        class_stats[class_id]['count'] += 1
                        class_stats[class_id]['total_confidence'] += confidence
                        class_stats[class_id]['confidences'].append(confidence)
                    
                    # 计算平均置信度
                    for class_id, stats in class_stats.items():
                        if stats['count'] > 0:
                            stats['avg_confidence'] = stats['total_confidence'] / stats['count']
                            stats['min_confidence'] = min(stats['confidences'])
                            stats['max_confidence'] = max(stats['confidences'])
                    
                    # 写入整体统计
                    f.write("整体统计:\n")
                    for class_id in sorted(class_stats.keys()):
                        stats = class_stats[class_id]
                        f.write(f"{stats['class_name']} (ID:{class_id}):\n")
                        f.write(f"  数量: {stats['count']} 张\n")
                        f.write(f"  平均置信度: {stats['avg_confidence']:.3f}\n")
                        f.write(f"  最小置信度: {stats['min_confidence']:.3f}\n")
                        f.write(f"  最大置信度: {stats['max_confidence']:.3f}\n")
                        f.write(f"  占比: {stats['count'] / len(all_results) * 100:.2f}%\n\n")
                    
                    # 写入详细结果
                    f.write("详细结果:\n")
                    for r in all_results:
                        f.write(f"{r['file_name']} - ")
                        f.write(f"类别: {r.get('class_name', '')} (ID: {r.get('class_id', '')}), ")
                        f.write(f"置信度: {r.get('confidence', 0):.3f}\n")
            
            print(f"汇总报告已保存: {summary_path}")
            
        except Exception as e:
            print(f"生成汇总报告失败: {str(e)}")


if __name__ == "__main__":
    # 测试代码
    model_path = "../weights/JinXiang_BALL/20251110_class/fullmodel_best.onnx"
    input_folder = "path/to/input/folder"
    output_folder = "path/to/output/folder"
    device = "cpu"
    
    service = BallLevelService(model_path=model_path, device=device)
    result = service.analyze_folder(input_folder=input_folder, output_folder=output_folder)
    
    if result['success']:
        print(f"分析成功！处理了 {result['processed_images']} 个文件")
    else:
        print(f"分析失败: {result['message']}")

